package com.wolfking.sharding.config;

import com.google.common.collect.ImmutableMap;
import io.micrometer.core.instrument.util.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.annotation.EnableTransactionManagement;

import javax.sql.DataSource;
import java.util.Objects;

@Configuration
@EnableTransactionManagement
public class DataSourceConfig {

    private static final ThreadLocal<String> contextHolder = new ThreadLocal<>();

    public static void setDataSourceType(String dataSourceType) {
        if (StringUtils.isNotBlank(dataSourceType)) {
            contextHolder.set(dataSourceType);
        }
    }

    public static String getDataSourceType() {
        return contextHolder.get();
    }

    public static void clearDataSourceType() {
        contextHolder.remove();
    }

    @Autowired
    @Qualifier("shardingSphereDataSource")
    private DataSource shardingMaster;

    @Autowired(required = false)
    @Qualifier("shardingSphereDataSourceSlave")
    private DataSource shardingSlave;

//    @Bean
//    @ConfigurationProperties(prefix = "spring.datasource.druid")
//    public DruidDataSource druidDataSource() {
//        return new DruidDataSource();
//    }


    @Bean
    @Primary
    public DataSource dataSource() {
        AbstractRoutingDataSource abstractRoutingDataSource = new AbstractRoutingDataSource() {
            @Override
            protected Object determineCurrentLookupKey() {
                return contextHolder.get();
            }
        };
        abstractRoutingDataSource.setDefaultTargetDataSource(shardingMaster);
        abstractRoutingDataSource.setTargetDataSources(ImmutableMap.of("master", shardingMaster, "slave",
                Objects.isNull(shardingSlave) ? shardingMaster : shardingSlave));
        return abstractRoutingDataSource;
    }


    @Bean
    public PlatformTransactionManager txManager(final DataSource dataSource) {
        return new DataSourceTransactionManager(dataSource);
    }
}
